from astropy.time import Time
from astropy import units as u
import numpy as np
import matplotlib.pyplot as plt
import cartopy.crs as ccrs
import cartopy.feature as cfeature
from cartopy.mpl.gridliner import LONGITUDE_FORMATTER, LATITUDE_FORMATTER
import os
import imageio.v2 as imageio
from joblib import Parallel, delayed

from pysagin.spacebase import Constellation, Satellite
from pysagin.groundbase import GroundStation
from pysagin.airbase import UavBase, Trajectory
from pysagin.device import RadioFixedAccess, RadioRandomAccess, Radio

# region 初始化全局参数
print("INFO: 初始化全局参数")

start_time = Time("2021-07-22 00:00:00")

# endregion

# region 创建 Kuiper 星座
print("INFO: 创建 Kuiper 星座")

constellation = Constellation(
    name="S",# Kuiper 4轨道星座
    type="Walker",
    T=40,  # 总卫星数量
    P=4,   # 轨道面数量
    F=0.5, # 相位因子
    a=6371+590,  # 半长轴（单位：km）
    e=0,  # 偏心率
    inc=53,  # 倾角（单位：度）
    init_time=start_time
)

constellation.update_all(start_time)
sats_gps = constellation.get_all_current_gps()
sats_lon_list = [sat_gps[0] for sat_gps in sats_gps]
sats_lat_list = [sat_gps[1] for sat_gps in sats_gps]
sats_alt_list = [sat_gps[2] for sat_gps in sats_gps]

# endregion

# region 创建地面站
print("INFO: 创建地面站")

station_1 = GroundStation(
    name="G1",# 河南省郑州市中原区须水镇
    init_time=start_time,
    init_gps=[113.51071, 34.755339, 0.0]
)

station_2 = GroundStation(
    name="G2",# 河南省郑州市荥阳市王村镇
    init_time=start_time,
    init_gps=[113.283777, 34.884903, 0.0]
)

station_3 = GroundStation(
    name="G3",# 河南省郑州市巩义市小关镇
    init_time=start_time,
    init_gps=[113.177662, 34.712468, 0.0]
)

stations_gps = [station_1.current_gps, station_2.current_gps, station_3.current_gps]
stations_lon_list = [station_gps[0] for station_gps in stations_gps]
stations_lat_list = [station_gps[1] for station_gps in stations_gps]
stations_alt_list = [station_gps[2] for station_gps in stations_gps]

# endregion

# region 创建无人机
print("INFO: 创建无人机")

uav_1_trajectory = Trajectory(name="U1Tr",avg_speed=(100*u.km/u.h).to(u.m/u.s).value)
uav_1_trajectory.add_straight(start=(105.883722,26.262808,0.0), end=(113.24587,34.729794,4000))
uav_1_trajectory.add_figure8(center=(113.24587,34.729794,4000), size=2000, loops=10)
uav_1_trajectory.add_straight(start=(113.24587,34.729794,4000), end=(105.883722,26.262808,0.0))
uav_1 = UavBase(name="U1",init_time=start_time,init_gps=[105.883722,26.262808,0.0])
uav_1.set_trajectory(trajectory=uav_1_trajectory, departure_time=start_time + 8 * u.hour)

uav_2_trajectory = Trajectory(name="U2Tr",avg_speed=(100*u.km/u.h).to(u.m/u.s).value)
uav_2_trajectory.add_straight(start=(105.883722,26.262808,0.0), end=(113.24587,34.729794,4000))
uav_2_trajectory.add_figure8(center=(113.24587,34.729794,4000), size=2000, loops=10)
uav_2_trajectory.add_straight(start=(113.24587,34.729794,4000), end=(105.883722,26.262808,0.0))
uav_2 = UavBase(name="U2",init_time=start_time,init_gps=[105.883722,26.262808,0.0])
uav_2.set_trajectory(trajectory=uav_2_trajectory, departure_time=start_time + 9 * u.hour)

uav_3_trajectory = Trajectory(name="U3Tr",avg_speed=(150*u.km/u.h).to(u.m/u.s).value)
uav_3_trajectory.add_straight(start=(116.409275,39.51952,0.0), end=(113.218588,34.800137,3500))
uav_3_trajectory.add_figure8(center=(113.218588,34.800137,3500), size=2000, loops=10)
uav_3_trajectory.add_straight(start=(113.218588,34.800137,3500), end=(116.409275,39.51952,0.0))
uav_3 = UavBase(name="U3",init_time=start_time,init_gps=[116.409275,39.51952,0.0])
uav_3.set_trajectory(trajectory=uav_3_trajectory, departure_time=start_time + 8 * u.hour)

uav_4_trajectory = Trajectory(name="U4Tr",avg_speed=(150*u.km/u.h).to(u.m/u.s).value)
uav_4_trajectory.add_straight(start=(116.409275,39.51952,0.0), end=(113.315836,34.795338,3500))
uav_4_trajectory.add_figure8(center=(113.315836,34.795338,3500), size=2000, loops=10)
uav_4_trajectory.add_straight(start=(113.315836,34.795338,3500), end=(116.409275,39.51952,0.0))
uav_4 = UavBase(name="U4",init_time=start_time,init_gps=[116.409275,39.51952,0.0])
uav_4.set_trajectory(trajectory=uav_4_trajectory, departure_time=start_time + 9 * u.hour)

uav_5_trajectory = Trajectory(name="U5Tr",avg_speed=(150*u.km/u.h).to(u.m/u.s).value)
uav_5_trajectory.add_straight(start=(116.409275,39.51952,0.0), end=(113.391887,34.767901,3500))
uav_5_trajectory.add_figure8(center=(113.391887,34.767901,3500), size=2000, loops=10)
uav_5_trajectory.add_straight(start=(113.391887,34.767901,3500), end=(116.409275,39.51952,0.0))
uav_5 = UavBase(name="U5",init_time=start_time,init_gps=[116.409275,39.51952,0.0])
uav_5.set_trajectory(trajectory=uav_5_trajectory, departure_time=start_time + 10 * u.hour)

uavs_gps = [uav_1.current_gps, uav_2.current_gps, uav_3.current_gps, 
            uav_4.current_gps, uav_5.current_gps]
uavs_lon_list = [uav_gps[0] for uav_gps in uavs_gps]
uavs_lat_list = [uav_gps[1] for uav_gps in uavs_gps]
uavs_alt_list = [uav_gps[2] for uav_gps in uavs_gps]

# endregion

# region 创建无线电台
print("INFO: 创建无线电台")

constellation.create_inter_satellite_links(Device_Radio=RadioFixedAccess, range=36000)

constellation_radio_random = []
for sat in constellation.satellites:
    sat_radio_random = RadioRandomAccess(name=f"{sat.name}D1", platform=sat, range=1000,
                                          max_access_num=-1)
    sat.devices.append(sat_radio_random)
    constellation_radio_random.append(sat_radio_random)


sation_radio_1 = RadioRandomAccess(name=f"{station_1.name}D0", platform=station_1, 
                                   range=1000, max_access_num=-1)
station_1.devices.append(sation_radio_1)
sation_radio_2 = RadioRandomAccess(name=f"{station_2.name}D0", platform=station_2, 
                                   range=1000, max_access_num=-1)
station_2.devices.append(sation_radio_2)
sation_radio_3 = RadioRandomAccess(name=f"{station_3.name}D0", platform=station_3, 
                                   range=1000, max_access_num=-1)
station_3.devices.append(sation_radio_3)


uav_radio_1 = RadioRandomAccess(name=f"{uav_1.name}D0", platform=uav_1, range=1000, max_access_num=-1)
uav_1.devices.append(uav_radio_1)
uav_radio_2 = RadioRandomAccess(name=f"{uav_2.name}D0", platform=uav_2, range=1000, max_access_num=-1)
uav_2.devices.append(uav_radio_2)
uav_radio_3 = RadioRandomAccess(name=f"{uav_3.name}D0", platform=uav_3, range=1000, max_access_num=-1)
uav_3.devices.append(uav_radio_3)
uav_radio_4 = RadioRandomAccess(name=f"{uav_4.name}D0", platform=uav_4, range=1000, max_access_num=-1)
uav_4.devices.append(uav_radio_4)
uav_radio_5 = RadioRandomAccess(name=f"{uav_5.name}D0", platform=uav_5, range=1000, max_access_num=-1)
uav_5.devices.append(uav_radio_5)

for sat_radio_random in constellation_radio_random:
    sat_radio_random.add_pre_access_radios([sation_radio_1, sation_radio_2, sation_radio_3]+
                                           [uav_radio_1, uav_radio_2, uav_radio_3, uav_radio_4, uav_radio_5])
sation_radio_1.add_pre_access_radios([uav_radio_1, uav_radio_2] + constellation_radio_random)
sation_radio_2.add_pre_access_radios([uav_radio_1, uav_radio_2] + constellation_radio_random)
sation_radio_3.add_pre_access_radios([uav_radio_1, uav_radio_2] + constellation_radio_random)
uav_radio_1.add_pre_access_radios([sation_radio_1, sation_radio_2, sation_radio_3] + constellation_radio_random)
uav_radio_2.add_pre_access_radios([sation_radio_1, sation_radio_2, sation_radio_3] + constellation_radio_random)
uav_radio_3.add_pre_access_radios([sation_radio_1, sation_radio_2, sation_radio_3] + constellation_radio_random)
uav_radio_4.add_pre_access_radios([sation_radio_1, sation_radio_2, sation_radio_3] + constellation_radio_random)
uav_radio_5.add_pre_access_radios([sation_radio_1, sation_radio_2, sation_radio_3] + constellation_radio_random)


ALL_DEVICES:list[Radio] = []
for sat in constellation.satellites:
    ALL_DEVICES.extend(sat.devices)
ALL_DEVICES.extend([sation_radio_1, sation_radio_2, sation_radio_3, 
                    uav_radio_1, uav_radio_2, uav_radio_3, uav_radio_4, uav_radio_5])

# endregion

# region 生成动画帧

print("INFO: 配置动画参数")
duration = 24 * u.hour  # 总时长
interval = 10 * u.min  # 时间间隔
time_steps = np.arange(0, duration.to(u.s).value, interval.to(u.s).value) * u.s
time_series = start_time + time_steps
total_frames = len(time_series)

def generate_frame(i, frame_time):
    """生成单帧图像并保存"""
    # 更新所有对象位置
    constellation.update_all(frame_time)
    sats_gps = constellation.get_all_current_gps()
    sats_lon_list = [sat_gps[0] for sat_gps in sats_gps]
    sats_lat_list = [sat_gps[1] for sat_gps in sats_gps]
    sats_alt_list = [sat_gps[2] for sat_gps in sats_gps]
    
    uav_1.update(frame_time)
    uav_2.update(frame_time)
    uav_3.update(frame_time)
    uav_4.update(frame_time)
    uav_5.update(frame_time)
    uavs_gps = [uav_1.current_gps, uav_2.current_gps, uav_3.current_gps, 
                uav_4.current_gps, uav_5.current_gps]
    uavs_lon_list = [uav_gps[0] for uav_gps in uavs_gps]
    uavs_lat_list = [uav_gps[1] for uav_gps in uavs_gps]
    uavs_alt_list = [uav_gps[2] for uav_gps in uavs_gps]
    
    # 创建地图
    import warnings
    warnings.filterwarnings("ignore", message="facecolor will have no effect as it has been defined as \"never\".")
    
    fig = plt.figure(figsize=(6, 4))
    ax = fig.add_subplot(1, 1, 1, projection=ccrs.PlateCarree())
    ax.set_extent([-180, 180, -90, 90], crs=ccrs.PlateCarree())
    
    # 添加地图特征
    ax.add_feature(cfeature.LAND, facecolor='grey', edgecolor='none', alpha=0.7)
    ax.add_feature(cfeature.OCEAN, facecolor='#A6CAE0', edgecolor='none')
    ax.add_feature(cfeature.COASTLINE, linewidth=0.1, color='white')
    
    # 绘制网格线
    gl = ax.gridlines(crs=ccrs.PlateCarree(), draw_labels=True,
                      linewidth=0.5, color='gray', alpha=0.5, linestyle='--')
    gl.top_labels = False
    gl.right_labels = False
    gl.xformatter = LONGITUDE_FORMATTER
    gl.yformatter = LATITUDE_FORMATTER
    
    # 绘制点
    ax.scatter(sats_lon_list, sats_lat_list, s=20, alpha=0.7, c="green", edgecolors="black", label="Satellites", transform=ccrs.PlateCarree())
    ax.scatter(stations_lon_list, stations_lat_list, s=20, alpha=0.7, c="blue", edgecolors="black", label="Ground Stations", transform=ccrs.PlateCarree())
    ax.scatter(uavs_lon_list, uavs_lat_list, s=20, alpha=0.7, c="red", edgecolors="black", label="UAVs", transform=ccrs.PlateCarree())
    
    # 绘制链路
    for dev_send in ALL_DEVICES:
        base_send = dev_send.platform
        radio_recvs = dev_send.get_current_access_radios()
        for radio_recv in radio_recvs:
            base_recv = radio_recv.platform
            try:
                ax.plot(
                    [base_send.current_gps[0], base_recv.current_gps[0]],
                    [base_send.current_gps[1], base_recv.current_gps[1]],
                    color="black", linewidth=0.1, transform=ccrs.Geodetic()
                )
            except:
                continue
    
    # ax.legend(loc='upper right', fontsize='small')
    plt.tight_layout()
    # 保存帧（按序号命名确保顺序）
    frame_path = f"frames/frame_{i:04d}.png"
    plt.savefig(frame_path, dpi=150)
    plt.close(fig)  # 关闭图像释放内存
    return frame_path

print("INFO: 并行生成动画帧...")

# 确保输出目录存在
os.makedirs("frames", exist_ok=True)
# 并行生成所有帧（n_jobs=-1表示使用所有CPU核心）
with Parallel(n_jobs=-1, verbose=10) as parallel:
    # 为每个帧分配索引和时间
    frame_paths = parallel(delayed(generate_frame)(i, t) for i, t in enumerate(time_series))

# endregion

# region 拼接帧为GIF
print("INFO: 拼接帧为GIF...")
# 按顺序读取所有帧
frames = []
for path in sorted(frame_paths):  # 确保按序号排序
    frames.append(imageio.imread(path))

# 保存为GIF
imageio.mimsave('SAGIN.gif', frames, fps=10, loop=0)
print("INFO: GIF生成完成！")

# 删除临时帧文件
import shutil
shutil.rmtree("frames")

# endregion
